package com.corn.zk.lock;

import lombok.SneakyThrows;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

/**
 * @author : Jim Wu
 * @version 1.0
 * @function :
 * @since : 2022/7/12 15:23
 */

public class DistributedLock {

    private final String zkServers;

    private final int sessionTimeout = 20000;

    // 连接建立完成latch
    private final CountDownLatch awaitZkConnectedLatch = new CountDownLatch(1);

    // 加锁等待latch
    private final CountDownLatch awaitLockRelease = new CountDownLatch(1);

    private String currentLockNode;

    private ZooKeeper zk;

    public DistributedLock(String zkServers) throws IOException, InterruptedException, KeeperException {
        this.zkServers = zkServers;
        this.init();
    }

    private void init() throws IOException, InterruptedException, KeeperException {
        // 连接zk
        zk = new ZooKeeper(zkServers, sessionTimeout, new Watcher() {
            @SneakyThrows
            @Override
            public void process(WatchedEvent event) {
                if (event.getState() == Event.KeeperState.SyncConnected) {
                    awaitZkConnectedLatch.countDown();
                }
            }
        });
        // 保证zk连接建立完成
        if (zk.getState() == ZooKeeper.States.CONNECTING) {
            awaitZkConnectedLatch.await();
        }
        // 创建一个/lock的根节点
        Stat exists = zk.exists("/lock", false);
        if (exists == null) {
            zk.create("/lock", null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }

    }

    public void lock() throws KeeperException, InterruptedException {
        // 创建zk带序号节点 并判断节点序号是不是最小的,如果是最小的说明当前应该获取锁,反之则挂起
        // 获取线程Id存入zk中
        ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
        buffer.putLong(Thread.currentThread().getId());
        // 创建ZK节点
        currentLockNode = zk.create("/lock/seq-", buffer.array(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        List<String> children = zk.getChildren("/lock", new Watcher() {
            @SneakyThrows
            @Override
            public void process(WatchedEvent event) {
                if (event.getType() == Event.EventType.NodeChildrenChanged) {
                    List<String> zkChildren = zk.getChildren("/lock", false);
                    // 获取/lock下面所有子节点并排序 如果是第一个表示当前应该获取锁 此时在事件监听中已经被await了,countDown释放阻塞
                    if (zkChildren.stream().map(child -> "/lock/" + child).sorted().collect(Collectors.toList()).indexOf(currentLockNode) <= 0) {
                        awaitLockRelease.countDown();
                    }
                }
            }
        });
        // 获取/lock下面所有子节点并排序 如果是第一个表示当前应该获取锁
        if (children.stream().map(child -> "/lock/" + child).sorted().collect(Collectors.toList()).indexOf(currentLockNode) <= 0) {
            return;
        } else {
            awaitLockRelease.await();
        }
    }

    public void unlock() throws KeeperException, InterruptedException {
        byte[] data = zk.getData(this.currentLockNode, false, null);
        if (Thread.currentThread().getId() == ByteBuffer.wrap(data).getLong()) {
            // 删除节点
            zk.delete(this.currentLockNode, -1);
        }
    }
}
